﻿using System.Collections.Generic;
using System.Linq;
using Algoritmia.DataStructures.Maps;
using Algoritmia.StateSpaces;
using Algoritmia.Utils;
using Lists = Algoritmia.DataStructures.Lists;
using Sets = Algoritmia.DataStructures.Sets;

namespace Algoritmia.Problems.ExactCover
{
    namespace DLX
    {
        #region Column2 derived class

        public class Column2 : Column
        {
            private readonly bool _primaryColumn;

            public Column2(bool primaryColumn, Cell left = null, Cell right = null, int s = 0)
                : base(left, right, s)
            {
                _primaryColumn = primaryColumn;
            }

            public bool IsPrimary
            {
                get { return _primaryColumn; }
            }
        }

        #endregion

        #region State2 class
        public class State2 : Cell
        {
            public Sets.ISet<int> Selected;
            public IMap<Cell, int> RowIndex;

            public State2(IEnumerable<IEnumerable<int>> ones, Lists.IList<int> selected, int n, int m, Sets.ISet<int> primaryColumns)
                : base(null, null, null, null)
            {
                RowIndex = AlgoritmiaFactory.Maps<Cell, int>.Default;
                Selected = AlgoritmiaFactory.Sets<int>.BuildDefault(selected);
                var list = ones.ToList();
                BuildColumnHeaders(m, primaryColumns);
                for (int row = 0; row < n; row++)
                {
                    Cell currentCol = this;
                    Cell first = null;
                    var currentList = list[row].ToList();
                    for (int col = 0; col < m; col++)
                    {
                        currentCol = currentCol.R;
                        if (currentList.Contains(col))
                        {
                            var column = (Column)currentCol;
                            column.S++;
                            var one = new One(currentCol.U, currentCol, column);
                            RowIndex[one] = row;
                            if (first == null)
                            {
                                first = one;
                            }
                            else
                            {
                                one.L = first.L;
                                one.R = first;
                                one.L.R = first.L = one;
                            }
                            currentCol.U = one.U.D = one;
                        }
                    }
                }

            }

            private void BuildColumnHeaders(int m, Sets.ISet<int> primaryColumns)
            {
                Cell ptr = this;
                for (int i = 0; i < m; i++)
                {
                    ptr.R = new Column2(primaryColumns.Contains(i));
                    ptr.R.L = ptr;
                    ptr = ptr.R;
                }
                ptr.R = this;
                L = ptr;
            }


        }
    }
    #endregion

    public class DLX2StateSpace : IReversibleForwardStateSpace<DLX.State2,DLX.Cell>
    {
        private readonly List<IEnumerable<int>> _ones;
        private readonly Lists.IList<int> _selected;
        private readonly Sets.ISet<int> _primaryColumns;
        private readonly int _n;
        private readonly int _m;

        public DLX2StateSpace(IEnumerable<IEnumerable<int>> ones,  IEnumerable<int> primaryColumns = null, Lists.IList<int> selected = null)
        {
            var onesList = ones.ToList();
            _primaryColumns = AlgoritmiaFactory.Sets<int>.BuildDefault(primaryColumns);
            _ones = onesList;
            _selected = selected ?? new int[] { }.ToAlgoritmiaList();
            _n = onesList.Count;
            _m = (from seq in onesList
                  select seq.Max()).Max() + 1;
        }

        #region Implementation of IForwardStateSpace<DLX.State2,DLX.Cell>

        public DLX.State2 Decide(DLX.State2 state, DLX.Cell cellInRow)
        {
            state.Selected.Add(state.RowIndex[cellInRow]);
            foreach (DLX.One one in cellInRow.Right())
            {
                one.Head.Cover();
            }
            return state;
        }

        public IEnumerable<DLX.Cell> Decisions(DLX.State2 state)
        {
            var size = _n + 1;
            DLX.Column2 bestCol = null;

            foreach(DLX.Column2 col in state.Right())
            {
                if ( col.IsPrimary && col.S < size)
                {
                    size = col.S;
                    bestCol = col;
                    if ( size == 1) break;
                }
            }

            bestCol.Cover();

            foreach(DLX.One cellInRow in bestCol.Down())
            {
                yield return cellInRow;
            }
            
            bestCol.UnCover();

        }

        public IEnumerable<DLX.State2> InitialStates
        {
            get { yield return new DLX.State2(_ones, _selected, _n, _m, _primaryColumns); }
        }

        public bool IsFinal(DLX.State2 state)
        {
            return state.Right().All(col => !((DLX.Column2)col).IsPrimary);
        }

        #endregion

        #region Implementation of IReversibleForwardStateSpace<DLX.State2,DLX.Cell>

        public DLX.State2 Undo(DLX.State2 state, DLX.Cell cellInRow)
        {
            state.Selected.Remove(state.RowIndex[cellInRow]);
            foreach (DLX.One one in cellInRow.Left())
            {
                one.Head.UnCover();
            }
            return state;
        }

        #endregion
    }
}
